from tensorflow.keras.models import load_model

from bert4keras.tokenizers import Tokenizer,load_vocab
from bert4keras.snippets import sequence_padding

#因为需要加载custom_object
from bert4keras.layers import *

from utils.dataloader import LCQMCLoader


from utils.bert_info import BertInfo

max_len = 128

model_saved_path = './model_saved/best_model'

bert_info_obj = BertInfo()

token_dict,keep_tokens = load_vocab(
    dict_path=bert_info_obj.dict_path,
    simplified=True,
    startswith=['[PAD]', '[UNK]', '[CLS]', '[SEP]'],
)
tokenizer = Tokenizer(token_dict,do_lower_case=True)

model = load_model(model_saved_path)

lcqmc_dataset = LCQMCLoader()
train_data = lcqmc_dataset.get_train_data()
valid_data = lcqmc_dataset.get_valid_data()
test_data = lcqmc_dataset.get_test_data()

data = valid_data
a_token_ids, b_token_ids, labels = [], [], []
texts = []

for d in data:
    token_ids = tokenizer.encode(d[0],maxlen=max_len)[0]
    a_token_ids.append(token_ids)
    token_ids = tokenizer.encode(d[1],maxlen=max_len)[0]
    b_token_ids.append(token_ids)
    labels.append(d[2])
    texts.extend(d[:2])
a_token_ids  = sequence_padding(a_token_ids)
b_token_ids = sequence_padding(b_token_ids)

a_vecs = model.predict([a_token_ids,np.zeros_like(a_token_ids)],verbose=True)
b_vecs = model.predict([b_token_ids,np.zeros_like(b_token_ids)],verbose=True)
labels = np.array(labels)

a_vecs = a_vecs / (a_vecs ** 2).sum(axis=1,keepdims=True) ** 0.5
b_vecs = b_vecs / (b_vecs ** 2).sum(axis=1,keepdims=True) ** 0.5
sims = (a_vecs * b_vecs).sum(axis=1)

# 以0.9为阈值，acc为79.82%
print('acc:', ((sims > 0.9) == labels.astype('bool')).mean())


# 测试全量检索能力
vecs = np.concatenate([a_vecs,b_vecs],axis=1).reshape((-1,768))

def most_similar(text,topn=10):
    token_ids,segment_ids = tokenizer.encode(text,maxlen=max_len)
    vec = model.predict([[token_ids],[segment_ids]])[0]
    vec /= (vec **2).sum()**0.5
    sims = np.dot(vecs,vec)
    return [(texts[i],sims[i]) for i in sims.argsort()[::-1][:topn]]

if __name__ == '__main__':
    most_similar(u'怎么开初婚未育证明', 20)